// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <include/ocr_rec.h>
#include <cstdio>
#include <locale.h>
// 需包含locale、string头文件、使用setlocale函数。

// string转wstring
std::wstring StringToWstring(const std::string str)
{
    unsigned len = str.size() * 2; // 预留字节数
    setlocale(LC_CTYPE, "");       //必须调用此函数
    wchar_t* p = new wchar_t[len]; // 申请一段内存存放转换后的字符串
    mbstowcs(p, str.c_str(), len); // 转换
    std::wstring str1(p);
    delete[] p;// 释放申请的内存
    return str1;
}

// wstring转string
std::string WstringToString(const std::wstring str)
{
    unsigned len = str.size() * 4;
    setlocale(LC_CTYPE, "");
    char* p = new char[len];
    wcstombs(p, str.c_str(), len);
    std::string str1(p);
    delete[] p;
    return str1;
}

namespace PaddleOCR {


//void CRNNRecognizer::Run(std::vector<cv::Mat> img_list,
//                         std::vector<std::string> &rec_texts,
//                         std::vector<float> &rec_text_scores,
//                         std::vector<double> &times) {
//    std::chrono::duration<float> preprocess_diff = std::chrono::duration<float>::zero();
//    std::chrono::duration<float> inference_diff = std::chrono::duration<float>::zero();
//    std::chrono::duration<float> postprocess_diff = std::chrono::duration<float>::zero();

//    int img_num = img_list.size();
//    std::vector<float> width_list;
//    for (int i = 0; i < img_num; i++) {
//        width_list.push_back(float(img_list[i].cols) / img_list[i].rows);
//    }
//    std::vector<int> indices = Utility::argsort(width_list);

//    for (int beg_img_no = 0; beg_img_no < img_num;
//         beg_img_no += this->rec_batch_num_) {
//        auto preprocess_start = std::chrono::steady_clock::now();
//        int end_img_no = std::min(img_num, beg_img_no + this->rec_batch_num_);
//        int batch_num = end_img_no - beg_img_no;
//        int imgH = this->rec_image_shape_[1];
//        int imgW = this->rec_image_shape_[2];
//        float max_wh_ratio = imgW * 1.0 / imgH;
//        for (int ino = beg_img_no; ino < end_img_no; ino++) {
//            int h = img_list[indices[ino]].rows;
//            int w = img_list[indices[ino]].cols;
//            float wh_ratio = w * 1.0 / h;
//            max_wh_ratio = std::max(max_wh_ratio, wh_ratio);
//        }

//        int batch_width = imgW;
//        std::vector<cv::Mat> norm_img_batch;
//        for (int ino = beg_img_no; ino < end_img_no; ino++) {
//            cv::Mat srcimg;
//            img_list[indices[ino]].copyTo(srcimg);
//            cv::Mat resize_img;
//            this->resize_op_.Run(srcimg, resize_img, max_wh_ratio,
//                                 this->use_tensorrt_, this->rec_image_shape_);
//            this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
//                                    this->is_scale_);
//            norm_img_batch.push_back(resize_img);
//            batch_width = std::max(resize_img.cols, batch_width);
//        }

//        std::vector<float> input(batch_num * 3 * imgH * batch_width, 0.0f);
//        this->permute_op_.Run(norm_img_batch, input.data());
//        auto preprocess_end = std::chrono::steady_clock::now();
//        preprocess_diff += preprocess_end - preprocess_start;
//        // Inference.
//        auto input_names = this->predictor_->GetInputNames();
//        auto input_t = this->predictor_->GetInputHandle(input_names[0]);
//        input_t->Reshape({batch_num, 3, imgH, batch_width});
//        auto inference_start = std::chrono::steady_clock::now();
//        input_t->CopyFromCpu(input.data());
//        this->predictor_->Run();

//        std::vector<float> predict_batch;
//        auto output_names = this->predictor_->GetOutputNames();
//        std::cout << output_names[0] << std::endl;
//        auto output_t = this->predictor_->GetOutputHandle(output_names[0]);
//        auto predict_shape = output_t->shape();

//        for(int i=0;i<predict_shape.size();i++)
//        {
//            std::cout << predict_shape[i] << std::endl;
//        }
//        int out_num = std::accumulate(predict_shape.begin(), predict_shape.end(), 1,
//                                      std::multiplies<int>());
//        std::cout << "size" << out_num << std::endl;
//        predict_batch.resize(out_num);
//        // predict_batch is the result of Last FC with softmax
//        output_t->CopyToCpu(predict_batch.data());
//        auto inference_end = std::chrono::steady_clock::now();
//        inference_diff += inference_end - inference_start;
//        // ctc decode
//        auto postprocess_start = std::chrono::steady_clock::now();
//        for (int m = 0; m < predict_shape[0]; m++) {
//            std::string str_res;
//            int argmax_idx;
//            int last_index = 0;
//            float score = 0.f;
//            int count = 0;
//            float max_value = 0.0f;



//            for (int n = 0; n < predict_shape[1]; n++) {
//                // get idx
//                argmax_idx = int(Utility::argmax(
//                                     &predict_batch[(m * predict_shape[1] + n) * predict_shape[2]],
//                        &predict_batch[(m * predict_shape[1] + n + 1) * predict_shape[2]]));
//                // get score
//                max_value = float(*std::max_element(
//                                      &predict_batch[(m * predict_shape[1] + n) * predict_shape[2]],
//                        &predict_batch[(m * predict_shape[1] + n + 1) * predict_shape[2]]));
//                std::cout << "id&score:" << argmax_idx << " "<< max_value << std::endl;
//                if (argmax_idx > 0 && (!(n > 0 && argmax_idx == last_index))) {
//                    score += max_value;
//                    count += 1;
//                    str_res += label_list_[argmax_idx];
//                }
//                last_index = argmax_idx;
//            }
//            score /= count;
//            if (std::isnan(score)) {
//                continue;
//            }
//            std::cout << str_res << std::endl;
//            rec_texts[indices[beg_img_no + m]] = str_res;
//            rec_text_scores[indices[beg_img_no + m]] = score;
//        }
//        auto postprocess_end = std::chrono::steady_clock::now();
//        postprocess_diff += postprocess_end - postprocess_start;
//    }
//    times.push_back(double(preprocess_diff.count() * 1000));
//    times.push_back(double(inference_diff.count() * 1000));
//    times.push_back(double(postprocess_diff.count() * 1000));
//}

//Group the decoded characters and record the corresponding decoded positions.

//Args:
//    text: the decoded text
//    selection: the bool array that identifies which columns of features are decoded as non-separated characters
//Returns:
//    word_list: list of the grouped words
//    word_col_list: list of decoding positions corresponding to each character in the grouped word
//    state_list: list of marker to identify the type of grouping words, including two types of grouping words:
//                - 'cn': continous chinese characters (e.g., 你好啊)
//                - 'en&num': continous english characters (e.g., hello), number (e.g., 123, 1.123), or mixed of them connected by '-' (e.g., VGG-16)
//                The remaining characters in text are treated as separators between groups (e.g., space, '(', ')', etc.).

void CRNNRecognizer::get_word_info(std::string text, std::vector<int> valid_col)
{


    std::vector<wchar_t> word_content;
    std::vector<int> word_col_content;
    std::vector<std::vector<wchar_t>> word_list;
    std::vector<std::vector<int>> word_col_list;

    std::vector<int> state_list;
    int state = CHAR_UNKNOW;
    std::wstring w_text = StringToWstring(text);
    for (int i=0; i<w_text.length(); i++)
    {
        int c_state = CHAR_UNKNOW;
        wchar_t ch = w_text.at(i);
        if( ch >= 0x4e00 && ch <= 0x9fff)
        {
            c_state = CHAR_CN;
        }
        else if((ch >= '0' && ch <='9') || (ch >= 'a' && ch <='z') || (ch >= 'A' && ch <='Z'))
        {
            c_state = CHAR_EN_NUM;
        }
        else
        {
            c_state = CHAR_SPLITTER;
        }


        if (ch == '.' && state == CHAR_EN_NUM && i < w_text.length()-1)
        {
            wchar_t ch_next = w_text.at(i-1);

            if((ch_next >= '0' && ch_next <='9'))
            {
                c_state = CHAR_EN_NUM;
            }
        }
        if (ch == '-' && state == CHAR_EN_NUM && i < w_text.length()-1)
        {
            c_state = CHAR_EN_NUM;
        }

        if (state == CHAR_UNKNOW)
        {
            state = c_state;
        }



        if(state != c_state)
        {
            if(word_content.size() != 0)
            {
                word_list.push_back(word_content);
                word_col_list.push_back(word_col_content);
                state_list.push_back(state);
                word_content.empty();
                word_col_content.empty();
            }
            state = c_state;
        }

        if (state != CHAR_SPLITTER)
        {
            word_content.push_back(ch);
            word_col_content.push_back(valid_col[i]);
        }
    }
    if(word_content.size() != 0)
    {
        word_list.push_back(word_content);
        word_col_list.push_back(word_col_content);
        state_list.push_back(state);
    }
    //    printf("word");
    //    return word_list, word_col_list, state_list
}
float getAverage(std::vector<float> &v) {
    float sum = 0;
    for (int i = 0; i < v.size(); i++) {
        sum += v[i];
    }
    return sum / v.size();
}
CharType charTypeofWchar(wchar_t ch)
{

    if( ch >= 0x3000)
    {
        return CHAR_CN;
        printf("%0x-cn, ",ch);
    }
    //                else if((ch >= '0' && ch <='9') || (ch >= 'a' && ch <='z') || (ch >= 'A' && ch <='Z'))
    //                {
    //                    c_state = CHAR_EN_NUM;
    //                }
    else
    {
        return CHAR_EN_NUM;
        printf("%0x-en, ", ch);
    }
}
void CRNNRecognizer::Run(std::vector<cv::Mat> img_list,
                         std::vector<std::string> &rec_texts,
                         std::vector<std::vector<float>> &rec_pos_list,
                         std::vector<std::vector<float>> &rec_width_list,
                         std::vector<float> &rec_text_scores,
                         std::vector<double> &times) {
    std::chrono::duration<float> preprocess_diff = std::chrono::duration<float>::zero();
    std::chrono::duration<float> inference_diff = std::chrono::duration<float>::zero();
    std::chrono::duration<float> postprocess_diff = std::chrono::duration<float>::zero();

    int img_num = img_list.size();
    std::vector<float> width_list;
    for (int i = 0; i < img_num; i++)
    {
        width_list.push_back(float(img_list[i].cols) / img_list[i].rows);
    }
    std::vector<int> indices = Utility::argsort(width_list);

    for (int beg_img_no = 0; beg_img_no < img_num;
         beg_img_no += this->rec_batch_num_)
    {
        auto preprocess_start = std::chrono::steady_clock::now();
        int end_img_no = std::min(img_num, beg_img_no + this->rec_batch_num_);
        int batch_num = end_img_no - beg_img_no;
        int imgH = this->rec_image_shape_[1];
        int imgW = this->rec_image_shape_[2];
        float max_wh_ratio = imgW * 1.0 / imgH;
        std::vector<float> wh_ratio_list;
        for (int ino = beg_img_no; ino < end_img_no; ino++)
        {
            int h = img_list[indices[ino]].rows;
            int w = img_list[indices[ino]].cols;
            float wh_ratio = w * 1.0 / h;
            wh_ratio_list.push_back(wh_ratio);
            max_wh_ratio = std::max(max_wh_ratio, wh_ratio);
        }

        int batch_width = imgW;
        std::vector<cv::Mat> norm_img_batch;
        for (int ino = beg_img_no; ino < end_img_no; ino++)
        {
            cv::Mat srcimg;
            img_list[indices[ino]].copyTo(srcimg);
            cv::Mat resize_img;
            this->resize_op_.Run(srcimg, resize_img, max_wh_ratio,
                                 this->use_tensorrt_, this->rec_image_shape_);
            this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                                    this->is_scale_);
            norm_img_batch.push_back(resize_img);
            batch_width = std::max(resize_img.cols, batch_width);
        }

        std::vector<float> input(batch_num * 3 * imgH * batch_width, 0.0f);
        this->permute_op_.Run(norm_img_batch, input.data());
        auto preprocess_end = std::chrono::steady_clock::now();
        preprocess_diff += preprocess_end - preprocess_start;
        // Inference.
        auto input_names = this->predictor_->GetInputNames();
        auto input_t = this->predictor_->GetInputHandle(input_names[0]);
        input_t->Reshape({batch_num, 3, imgH, batch_width});
        auto inference_start = std::chrono::steady_clock::now();
        input_t->CopyFromCpu(input.data());
        this->predictor_->Run();

        std::vector<float> predict_batch;
        auto output_names = this->predictor_->GetOutputNames();
        //        std::cout << output_names[0] << std::endl;
        auto output_t = this->predictor_->GetOutputHandle(output_names[0]);
        auto predict_shape = output_t->shape();

        int out_num = std::accumulate(predict_shape.begin(), predict_shape.end(), 1,
                                      std::multiplies<int>());
        //        std::cout << "size" << out_num << std::endl;
        predict_batch.resize(out_num);
        // predict_batch is the result of Last FC with softmax
        output_t->CopyToCpu(predict_batch.data());
        auto inference_end = std::chrono::steady_clock::now();
        inference_diff += inference_end - inference_start;
        // ctc decode
        auto postprocess_start = std::chrono::steady_clock::now();

        int ignored_tokens[1] = {0};
        printf("batch count: %d\n", predict_shape[0]);
        for (int m = 0; m < predict_shape[0]; m++)
        {
            float widthf = predict_shape[1]*(wh_ratio_list[m]/max_wh_ratio);
            std::string str_res;
            int argmax_idx;

            float score = 0.f;
            int count = 0;
            float max_value = 0.0f;
            int shape_count = predict_shape[1];
            bool *selection = (bool*)malloc(shape_count * sizeof(bool));

            if (selection == nullptr) {
                std::cout << "selection内存分配失败！" << std::endl;
                return ;
            }

            // 对分配的内存空间进行初始化（这里选择将每个元素都赋值为0）
            memset(selection, true, shape_count * sizeof(bool));

            int  *text_index = (int*)malloc(shape_count * sizeof(int));
            if (text_index == nullptr) {
                std::cout << "text_index内存分配失败！" << std::endl;
                return ;
            }
            float  *text_prob = (float*)malloc(shape_count * sizeof(float));
            if (text_prob == nullptr) {
                std::cout << "text_index内存分配失败！" << std::endl;
                return ;
            }



            for (int n = 0; n < predict_shape[1]; n++) {
                // get idx
                text_index[n] = int(Utility::argmax(
                                        &predict_batch[(m * predict_shape[1] + n) * predict_shape[2]],
                        &predict_batch[(m * predict_shape[1] + n + 1) * predict_shape[2]]));
                // get score
                text_prob[n] = float(*std::max_element(
                                         &predict_batch[(m * predict_shape[1] + n) * predict_shape[2]],
                        &predict_batch[(m * predict_shape[1] + n + 1) * predict_shape[2]]));
            }


            //remove duplicate
            for (int n = 1; n < shape_count; n++) {
                selection[n] = text_index[n] != text_index[n-1];
            }
            //ignored tokens
            for (int n = 0; n < shape_count; n++) {
                selection[n] &= (text_index[n] != 0);
            }

            std::vector<int> valid_col;
            for (int n = 0; n < shape_count; n++) {
                if(selection[n])
                {
                    str_res += label_list_[text_index[n]];
                    score += text_prob[n];
                    count += 1;
                    valid_col.push_back(n);
                }
            }
            score /= count;
            if (std::isnan(score)) {
                continue;
            }
            int state = CHAR_UNKNOW;
            std::wstring w_text = StringToWstring(str_res);
            int w = img_list[indices[beg_img_no + m]].cols;
            float cell_width = w/widthf;
            int c_state = CHAR_UNKNOW;
            int last_index = 0;
            std::vector<int> cn_index_list;
            std::vector<float> cn_width_list;
            std::vector<float> en_width_list;
            std::vector<float> pos_list(valid_col.size(), 0.0);
            std::vector<float> width_list(valid_col.size(), 0.0);
            bool changed = false;
            for (int i=0; i<w_text.length(); i++)
            {
                changed = false;
                c_state = charTypeofWchar(w_text.at(i));
                if(c_state == CHAR_CN)
                {
                    cn_index_list.push_back(i);
                }

                if(i+1<w_text.length())
                {
                    if(c_state == charTypeofWchar(w_text.at(i+1)))
                    {
                        continue;
                    }
                    changed = true;
                }
                if(c_state == CHAR_CN)
                {
                    if(i != last_index)
                    {
                        float char_seq_length = (valid_col.at(i) - valid_col.at(last_index) + 1) * cell_width;
                        float char_width = char_seq_length/(i-last_index);
                        cn_width_list.push_back(char_width);
                    }
                }
                else
                {
                    if(i != last_index)
                    {
                        float char_seq_length = (valid_col.at(i) - valid_col.at(last_index) + 1) * cell_width;
                        float char_width = char_seq_length/(i-last_index);
                        en_width_list.push_back(char_width);
                    }
//                    float start = (valid_col.at(lastIndex)-1) * cell_width;
//                    float end = (valid_col.at(i)+1) * cell_width;
//                    float ave = (end-start)/(i-lastIndex+1);
//                    for(int k=lastIndex, num=0;k<=i;k++, num++)
//                    {
//                        pos_list[k] = start+ave*num;
//                        width_list[k] = ave;
//                    }
//                    pos_list[i] = start;
//                    width_list[i] = end-start;
                }
                if(changed)
                {
                    last_index = i+1;
                }
            }
            float avg_cn_width = 0;
            if(cn_width_list.size() == 0)
            {
                avg_cn_width = w/w_text.length();
            }
            else
            {
                avg_cn_width = getAverage(cn_width_list);
            }
            float avg_en_width = 0;
            if(cn_width_list.size() == 0)
            {
                avg_en_width = w/w_text.length();
            }
            else
            {
                avg_en_width = getAverage(en_width_list);
            }
            last_index=0;
            int last_state = CHAR_UNKNOW;
            int cur_state = CHAR_UNKNOW;
            for (int i=0,j=0; i<w_text.length(); i++)
            {
                if(cn_index_list.size()>0 && cn_index_list[j]==i)
                {
                    j++;
                    cur_state = CHAR_CN;
                }
                else
                {
                    cur_state = CHAR_EN;
                }
                if(i>0)
                {
                    if(last_state != cur_state)
                    {
                        for(int k=last_index;k<=i-1;k++)
                        {
                            float center_x = (valid_col.at(k)+0.5)*cell_width;
                            float start=0,end=0;
                            if(last_state == CHAR_CN)
                            {
                                start = fmaxf(center_x - avg_cn_width/2, 0);
                                end = fmin(center_x + avg_cn_width/2, w);
                            }
                            else
                            {
                                start = fmaxf(center_x - avg_en_width/2, 0);
                                end = fmin(center_x + avg_en_width/2, w);
                            }
                            pos_list[k] = start;
                            width_list[k] = end-start;
                        }
                        last_index=i;
                    }
                }
                last_state = cur_state;
            }
            for(int k=last_index;k<=w_text.length()-1;k++)
            {
                float center_x = (valid_col.at(k)+0.5)*cell_width;
                float start=0,end=0;
                if(last_state == CHAR_CN)
                {
                    start = fmaxf(center_x - avg_cn_width/2, 0);
                    end = fmin(center_x + avg_cn_width/2, w);
                }
                else
                {
                    start = fmaxf(center_x - avg_en_width/2, 0);
                    end = fmin(center_x + avg_en_width/2, w);
                }
                pos_list[k] = start;
                width_list[k] = end-start;
            }

            rec_texts[indices[beg_img_no + m]] = str_res;
            rec_pos_list[indices[beg_img_no + m]] = pos_list;
            rec_width_list[indices[beg_img_no + m]] = width_list;
            rec_text_scores[indices[beg_img_no + m]] = score;

            // 释放内存空间
            free(selection);
            free(text_index);
            free(text_prob);
        }
        auto postprocess_end = std::chrono::steady_clock::now();
        postprocess_diff += postprocess_end - postprocess_start;
    }
    times.push_back(double(preprocess_diff.count() * 1000));
    times.push_back(double(inference_diff.count() * 1000));
    times.push_back(double(postprocess_diff.count() * 1000));
}

void CRNNRecognizer::LoadModel(const std::string &model_dir) {
    paddle_infer::Config config;
    config.SetModel(model_dir + "/inference.pdmodel",
                    model_dir + "/inference.pdiparams");
    std::cout << "In PP-OCRv3, default rec_img_h is 48,"
              << "if you use other model, you should set the param rec_img_h=32"
              << std::endl;
    if (this->use_gpu_) {
        config.EnableUseGpu(this->gpu_mem_, this->gpu_id_);
        if (this->use_tensorrt_) {
            auto precision = paddle_infer::Config::Precision::kFloat32;
            if (this->precision_ == "fp16") {
                precision = paddle_infer::Config::Precision::kHalf;
            }
            if (this->precision_ == "int8") {
                precision = paddle_infer::Config::Precision::kInt8;
            }
            if (!Utility::PathExists("./trt_rec_shape.txt")) {
                config.CollectShapeRangeInfo("./trt_rec_shape.txt");
            } else {
                config.EnableTunedTensorRtDynamicShape("./trt_rec_shape.txt", true);
            }
        }
    } else {
        config.DisableGpu();
        if (this->use_mkldnn_) {
            config.EnableMKLDNN();
            // cache 10 different shapes for mkldnn to avoid memory leak
            config.SetMkldnnCacheCapacity(10);
        }
        config.SetCpuMathLibraryNumThreads(this->cpu_math_library_num_threads_);
    }

    // get pass_builder object
    auto pass_builder = config.pass_builder();
    // delete "matmul_transpose_reshape_fuse_pass"
    pass_builder->DeletePass("matmul_transpose_reshape_fuse_pass");
    config.SwitchUseFeedFetchOps(false);
    // true for multiple input
    config.SwitchSpecifyInputNames(true);

    config.SwitchIrOptim(true);

    config.EnableMemoryOptim();
    //   config.DisableGlogInfo();

    this->predictor_ = paddle_infer::CreatePredictor(config);
}

} // namespace PaddleOCR
